
import numpy as np


def preprocess(text):
    """文本（语料库）预处理
    对一个非常小的文本数据（语料库）进行预处理。
    就是将文本分割为单词（分词），
    并将分割后的单词列表转化为单词ID列表。
    
    parameters
    -----------
    text: 输入文本

    returns
    ----------
    corpus:单词ID列表
    word_to_id: 单词到单词ID的字典
    id_to_word: 单词ID到单词的字典
    """
    text = text.lower()
    text = text.replace('.',' .')
    words = text.split()

    word_to_id,id_to_word = {},{}
    for word in words:
        if word not in word_to_id:
            new_id = len(word_to_id)
            word_to_id[word]=new_id
            id_to_word[new_id]=word 
    
    corpus = [word_to_id[word] for word in words]   

    return corpus,word_to_id,id_to_word 

# 分布式假设（distributional hypothesis):
#     某个单词的含义由它周围的单词形成
# 也就是说单词含义由他所在的上下文形成。
def create_co_matrix(corpus,vocab_size,window_size=1):
    """创建共现矩阵
    共现矩阵的每个行，就是对应单词的向量表示

    parameters
    ------------
    corpus: 语料库
    vocab_size: 语料库包含的单词个数
    window_size: 窗口大小，表示上下文统计单词的个数，当1时，就是只考虑左右各1

    returns
    ----------
    返回共现矩阵，其中每行表示语料库中一个单词的向量表示
    """
    corpus_size = len(corpus)
    co_matrix = np.zeros((vocab_size,vocab_size),dtype=np.int32)

    for idx,word_id in enumerate(corpus):
        for i in range(1,window_size+1):
            left_idx = idx - i 
            right_idx = idx + i 

            if left_idx >= 0:
                left_word_id = corpus[left_idx]
                co_matrix[word_id,left_word_id] += 1
            if right_idx < corpus_size:
                right_word_id = corpus[right_idx]
                co_matrix[word_id,right_word_id] += 1
    return co_matrix

# 单词相似度：可以用来衡量两个单词词义接近
# 相似度可以采用向量的余弦相似度
def cos_similarity(x,y,eps=1e-8):
    '''计算两个单词（词向量）的余弦相似度

    parameters
    -----------
    x: 单词向量
    y: 单词向量
    eps: 防止“除数为0”的微小值
    
    returns
    一个实数，代表两个向量的相似度
    '''
    # 1. 手动计算------------
    # 计算向量的模（欧几里得范数，L2范数）
    # x_norm = np.sqrt(np.sum(x**2))
    # y_norm = np.sqrt(np.sum(y**2))

    # 2. 使用np.linalg.norm计算
    x_norm = np.linalg.norm(x)
    y_norm = np.linalg.norm(y)

    # 单位化向量
    nx = x/(x_norm+eps)
    ny = y/(y_norm+eps)

    # 两个单位向量的内积，就是夹角余弦
    return np.dot(nx,ny)

# 相似单词排序：
#     当某个单词被作为查询词时，返回与这个查询词相似
# 的单词降序结果。
def most_similar(query,word_to_id,id_to_word,word_matrix,high=5):
    """相似单词查找
    
    parameters
    -----------
    query: 查询词
    word_to_id: 从单词到单词ID的字典
    id_to_word: 从单词ID到单词的字典
    word_matrix: 单词向量矩阵
    top: 显示多少个要关联的单词
    """
    
    if query not in word_to_id:
        print(f"{query} is not found")
        return 
    
    # 获取查询单词的词向量
    print(f"\n[query] {query}")
    query_id = word_to_id[query]
    query_vec = word_matrix[query_id]

    vocab_size = len(id_to_word)

    similarity = np.zeros(vocab_size)
    # 与所有其他单词的余弦相似度
    for i in range(vocab_size):
        similarity[i] =cos_similarity(word_matrix[i],query_vec)

    count = 0
    for i in (-1*similarity).argsort():
        if id_to_word[i] == query:
            continue
        print(' %s: %s' % (id_to_word[i], similarity[i]))

        count += 1
        if count>=high:
            return 

def ppmi(C,verbose=False,eps=1e-8):
    M = np.zeros_like(C,dtype=np.float32)
    N = np.sum(C) # 用于计算联合概率和边缘概率
    S = np.sum(C,axis=0) # 用于计算边缘概率

    total = C.shape[0]*C.shape[1]
    cnt = 0

    for i in range(C.shape[0]):
        for j in range(C.shape[1]):
            # pmi = np.log2(C[i,j]*N/(S[j]*S[i])+eps)
            # 避免运算中的警告，采用下面的运算
            pmi = np.log2(C[i,j]+ eps) + np.log2(N+ eps) - np.log2(S[j]+ eps) - np.log2(S[i]+ eps)

            M[i,j] = max(0,pmi)

            if verbose:
                cnt += 1
                if cnt % (total//100+1) == 0:
                    ratio = cnt/total
                    # print('%.1f%% done' % (ratio))
                    print(f"{ratio:.1%} done")
    return M

def create_contexts_target(corpus,window_size=1):
    target = corpus[window_size:-window_size]
    contexts = []

    for idx in range(window_size,len(corpus)-window_size):
        cs = []
        for t in range(-window_size, window_size+1):
            if t == 0:
                continue
            cs.append(corpus[idx+t])
        contexts.append(cs)
    
    return np.array(contexts), np.array(target)

def convert_one_hot(corpus,vocab_size):
    """将单词ID转换为one-hot-vector
    
    parameters
    ------------
    corpus: 单词ID列表（一维或二维数组）
    vocab_size: 词汇个数（决定向量长度）
    
    returns
    --------
    one-hot表示（二维或三维数组）
    """
    N = corpus.shape[0]

    if corpus.ndim == 1:
        one_hot = np.zeros((N,vocab_size),dtype=np.int32)
        for idx,word_id in enumerate(corpus):
            one_hot[idx,word_id] = 1

    elif corpus.ndim == 2:
        C = corpus.shape[1]
        one_hot = np.zeros((N,C,vocab_size),dtype=np.int32)
        for idx_0,word_ids in enumerate(corpus):
            for idx_1,word_id in enumerate(word_ids):
                one_hot[idx_0,idx_1,word_id] = 1
    
    return one_hot

def analogy(a, b, c, word_to_id, id_to_word, word_matrix, top=5, answer=None):
    for word in (a, b, c):
        if word not in word_to_id:
            print('%s is not found' % word)
            return

    print('\n[analogy] ' + a + ':' + b + ' = ' + c + ':?')
    a_vec, b_vec, c_vec = word_matrix[word_to_id[a]], word_matrix[word_to_id[b]], word_matrix[word_to_id[c]]
    query_vec = b_vec - a_vec + c_vec
    query_vec = normalize(query_vec)

    similarity = np.dot(word_matrix, query_vec)

    if answer is not None:
        print("==>" + answer + ":" + str(np.dot(word_matrix[word_to_id[answer]], query_vec)))

    count = 0
    for i in (-1 * similarity).argsort():
        if np.isnan(similarity[i]):
            continue
        if id_to_word[i] in (a, b, c):
            continue
        print(' {0}: {1}'.format(id_to_word[i], similarity[i]))

        count += 1
        if count >= top:
            return

def normalize(x):
    if x.ndim == 2:
        s = np.sqrt((x * x).sum(1))
        x /= s.reshape((s.shape[0], 1))
    elif x.ndim == 1:
        s = np.sqrt((x * x).sum())
        x /= s
    return x

# 梯度剪裁，通过限制梯度的最大范数来防止梯度爆炸
def clip_grads(grads,max_norm):
    total_norm = 0
    for grad in grads:
        total_norm += np.sum(grad**2)
    total_norm = np.sqrt(total_norm)

    rate = max_norm/(total_norm+1e-6)
    if rate < 1:
        for grad in grads:
            grad *= rate 



if __name__ == '__main__':
    text = "You say goodbye and I say hello."
    corpus,word_to_id,id_to_word = preprocess(text)
    # print(corpus)
    # C = create_co_matrix(corpus,len(word_to_id))
    # M = ppmi(C,verbose=False)
    # print(M)
    grads=np.array([0.9,1.2,2.0])
    clip_grads(grads,1.0)
